import json, os
from pathlib import Path
from data_cleaning import DataCleaningAgent
ROOT = Path(__file__).resolve().parents[2]
RAW_DIR = ROOT / "data" / "raw"
PROC_DIR = ROOT / "data" / "processed"
PROC_DIR.mkdir(parents=True, exist_ok=True)

def read_raw_samples():
    files = sorted(RAW_DIR.glob("*.txt"))
    if not files:
        # 内置两条示例，避免“无数据就跑不起来
        return [
            "地址：南昌市东湖区，北京华为技术有限公司；联系人：王五，电话：13800001234"
        ]
    samples = []
    for f in files:
        samples.append(f.read_text(encoding="utf-8"))
    return samples

if __name__ == "__main__":
    agent = DataCleaningAgent()
    samples = read_raw_samples()
    instruction = "请提取company(公司名称)、contact(联系人)、phone(电话)。返回JSON。"
    results = agent.batch_process(samples, instruction)
    # 写入两个文件，便于排查与复用
    with open(PROC_DIR / "cleaned.jsonl", "w", encoding="utf-8") as fw1, \
        open(PROC_DIR / "extracted.jsonl", "w", encoding="utf-8") as fw2:
        for r in results:
            fw1.write(json.dumps({"original": r["original"], "cleaned": r["cleaned"]},ensure_ascii=False) + "\n")
            fw2.write(json.dumps(r["extracted"], ensure_ascii=False) + "\n")
    print(f"已输出：{PROC_DIR/'cleaned.jsonl'} 与 {PROC_DIR/'extracted.jsonl'}")